#include "cSvsOpencv.h"

cSvsOpencv::cSvsOpencv(int aResH, int aResV, bool aColor, int aStereoAlgo, const char * aExtrinFileName, const char * aIntrinFileName)
{
	resH = aResH;
	resV = aResV;
	isColor = aColor;
	stereoAlg = aStereoAlgo;
	SADWindowSize = 0;
	numberOfDisparities = 0;

	if(aExtrinFileName != NULL)
	{
		extrinFileName = aExtrinFileName;
		intrinFileName = aIntrinFileName;
	}

	videoObject = getVideoObject();
	
	bpm = new cv::gpu::StereoBeliefPropagation(64, 8, 2, 25, 0.1f, 15, 1, CV_16S);
	bmGPU = new cv::gpu::StereoBM_GPU();//(0, 64, 9);
	csbpm = new cv::gpu::StereoConstantSpaceBP(32, 8, 4, 4);
}

cSvsOpencv::~cSvsOpencv()
{
	this->Stop();
	this->Close();
}

void cSvsOpencv::Close()
{
	if(opened)
	{
		videoObject->Close();
		opened = false;
	}
}

bool cSvsOpencv::Open()
{
	if(!opened)
	{
		bool ret = videoObject->Open();
		if(!ret)
			return false;
		opened = true;
	}

	return true;
}

void cSvsOpencv::Start()
{
	if(!started && opened)
	{

		if(isColor)
		{
			headerLeft = cvCreateImageHeader(cvSize(resH, resV), IPL_DEPTH_8U, 4);
			headerRight = cvCreateImageHeader(cvSize(resH, resV), IPL_DEPTH_8U, 4);
		}
		else
		{
			headerLeft = cvCreateImageHeader(cvSize(resH, resV), IPL_DEPTH_8U, 1);
			headerRight = cvCreateImageHeader(cvSize(resH, resV), IPL_DEPTH_8U, 1);
		}

		imgLeft = cvCreateImage(cvSize(resH,resV),IPL_DEPTH_8U,3);
		imgRight = cvCreateImage(cvSize(resH,resV),IPL_DEPTH_8U,3);

		videoObject->SetColor(true, true);
		videoObject->SetFrameDiv(1); //default
		videoObject->SetSize(resH, resV);
		this->initRetification();
		videoObject->Start();
		started = true;
	}
}

void cSvsOpencv::Stop()
{
	if(started)
	{
		cvReleaseImageHeader(&headerLeft);
		cvReleaseImageHeader(&headerLeft);
		cvReleaseImage(&imgLeft);
		cvReleaseImage(&imgRight);
		videoObject->Stop();
		started = false;
	}
}

IplImage* cSvsOpencv::calculateStereo()
{
	Mat img1(headerLeft), img2(headerRight);
	Size img_size = img1.size();

	Mat img1r, img2r;
	remap(img1, img1r, map11, map12, INTER_LINEAR);
	remap(img2, img2r, map21, map22, INTER_LINEAR);

	img1 = img1r;
	img2 = img2r;

	numberOfDisparities = numberOfDisparities > 0 ? numberOfDisparities : 32;

	bm.state->roi1 = roi1;
	bm.state->roi2 = roi2;
	bm.state->preFilterCap = 31;
	bm.state->SADWindowSize = SADWindowSize > 0 ? SADWindowSize : 9;
	bm.state->minDisparity = 0;
	bm.state->numberOfDisparities = numberOfDisparities;
	bm.state->textureThreshold = 10;
	bm.state->uniquenessRatio = 15;
	bm.state->speckleWindowSize = 100;
	bm.state->speckleRange = 32;
	bm.state->disp12MaxDiff = 1;

	sgbm.preFilterCap = 63;
	sgbm.SADWindowSize = SADWindowSize > 0 ? SADWindowSize : 3;

	int cn = img1.channels();

	sgbm.P1 = 8 * cn * sgbm.SADWindowSize * sgbm.SADWindowSize;
	sgbm.P2 = 32 * cn * sgbm.SADWindowSize * sgbm.SADWindowSize;
	sgbm.minDisparity = 0;
	sgbm.numberOfDisparities = numberOfDisparities;
	sgbm.uniquenessRatio = 10;
	sgbm.speckleWindowSize = bm.state->speckleWindowSize;
	sgbm.speckleRange = bm.state->speckleRange;
	sgbm.disp12MaxDiff = 1;
	sgbm.fullDP = stereoAlg == STEREO_HH;

	Mat disp, disp8;
	IplImage iplImg1 = img1;
	IplImage iplImg2 = img2;	
	gpu::GpuMat dispGPU;
	
	switch(stereoAlg)
	{
		case STEREO_BM: {bm(img1, img2, disp); break;}
		case STEREO_BP: {(*bpm)(gpu::GpuMat(img1), gpu::GpuMat(img2), dispGPU); break; }
		case STEREO_BMGPU: {(*bmGPU)(gpu::GpuMat(img1), gpu::GpuMat(img2), dispGPU); break; }
		case STEREO_CSBP: {(*csbpm)(gpu::GpuMat(img1), gpu::GpuMat(img2), dispGPU); break; }
		default: {sgbm(img1, img2, disp); break;}
	}

	if(stereoAlg == STEREO_BP || stereoAlg == STEREO_BMGPU || stereoAlg == STEREO_CSBP)
	{
		dispGPU.convertTo(dispGPU, img1.type());
		dispGPU.download(disp8);
	}
	else disp.convertTo(disp8, CV_8U, 255 / (numberOfDisparities * 16.));

	IplImage disp8old = disp8;
	static IplImage* dispout = cvCreateImage(cvSize(disp8old.width,disp8old.height),IPL_DEPTH_8U,3);  //memory leak
	cvConvertImage(&disp8old, dispout, CV_CVTIMG_SWAP_RB);

	return dispout;
}

IplImage* cSvsOpencv::getLeft()
{
	return imgLeft;
}

IplImage* cSvsOpencv::getRight()
{
	return imgRight;
}

void cSvsOpencv::setColor(bool color)
{
	isColor = color;
}

void cSvsOpencv::setExtrinsicsFile(const char* aExtrinFileName)
{
	extrinFileName = aExtrinFileName;

}

void cSvsOpencv::setIntrisicsFile(const char* aIntrinFileName)
{
	intrinFileName = aIntrinFileName;
}

void cSvsOpencv::setRes(int aResH, int aResV)
{
	resH = aResH;
	resV = aResV;
}

void cSvsOpencv::setStereoAlgo(int alg)
{
	stereoAlg = alg;
}
void cSvsOpencv::refresh()
{
	if(started && opened)
	{
		si = videoObject->GetImage(500);
		if(isColor)
		{
			headerLeft->imageData = (char*) si->color;
			headerRight->imageData = (char*) si->color_right;
		}
		else
		{
			headerLeft->imageData = (char*) si->left;
			headerRight->imageData = (char*) si->right;
		}

		if(isColor)
		{
			cvCvtColor(headerLeft, imgLeft, CV_BGR2RGB);
			cvCvtColor(headerRight, imgRight, CV_BGR2RGB);
		}
		else
		{
			cvConvertImage(headerLeft, imgLeft, CV_CVTIMG_SWAP_RB);
			cvConvertImage(headerRight, imgRight, CV_CVTIMG_SWAP_RB);
		}
	}
}

void cSvsOpencv::initRetification()
{
	Mat Q;
	Size img_size = cvSize(resH, resV);

	FileStorage fs(intrinFileName, CV_STORAGE_READ);
	if (!fs.isOpened())
	{
		printf("Failed to open file %s\n", intrinFileName.c_str());
		return;
	}

	Mat M1, D1, M2, D2;
	fs["M1"] >> M1;
	fs["D1"] >> D1;
	fs["M2"] >> M2;
	fs["D2"] >> D2;

	fs.open(extrinFileName, CV_STORAGE_READ);
	if (!fs.isOpened())
	{
		printf("Failed to open file %s\n", extrinFileName.c_str());
		return;
	}

	Mat R, T, R1, P1, R2, P2;
	fs["R"] >> R;
	fs["T"] >> T;

	stereoRectify(M1, D1, M2, D2, img_size, R, T, R1, R2, P1, P2, Q, -1, img_size, &roi1, &roi2);

	initUndistortRectifyMap(M1, D1, R1, P1, img_size, CV_16SC2, map11, map12);
	initUndistortRectifyMap(M2, D2, R2, P2, img_size, CV_16SC2, map21, map22);

}
void cSvsOpencv::setNumberOfDisparities(int nDisparities)
{
	numberOfDisparities = nDisparities;
}

void cSvsOpencv::setSADWindowSize(int wSize)
{
	SADWindowSize = wSize;
}
